package com.zz.ai.util;

import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.document.Document;

import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

public class CustomTokenTextSplitter extends TokenTextSplitter {

    private static final int CHUNK_SIZE = 1500;

    private static final int MIN_CHUNK_SIZE_CHARS = 1000;

    private static final int MIN_CHUNK_LENGTH_TO_EMBED = 5;

    private static final int MAX_NUM_CHUNKS = 10000;

    private static final boolean KEEP_SEPARATOR = true;

    // 继承父类参数配置
    public CustomTokenTextSplitter() {
        super(CHUNK_SIZE, MIN_CHUNK_SIZE_CHARS, MIN_CHUNK_LENGTH_TO_EMBED, MAX_NUM_CHUNKS,KEEP_SEPARATOR);
    }

    @Override
    public List<Document> split(Document document) {
        // 按 ### 分割原始文本
        String[] sections = document.getText().split("###");

        return Arrays.stream(sections)
                .filter(section -> !section.isBlank())
                .map(section -> {
                    // 对每个分割后的段落进行 Token 处理
                    List<Document> chunks = super.split(new  Document(section.trim()));
                    return chunks;
                })
                .flatMap(List::stream)
                .collect(Collectors.toList());
    }
}